# -*- coding: utf-8 -*-
"""
路径规划强化学习环境
"""
#
import gym
import math
import random
import numpy as np
from gym import spaces
from copy import deepcopy
from pathlib import Path
from collections import deque
from scipy.integrate import odeint,solve_ivp
from shapely import geometry as geo
from shapely.plotting import plot_polygon
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

__all__ = ["DynamicPathPlanning"]


# ----------------------------- ↓↓↓↓↓ 地图设置 ↓↓↓↓↓ ------------------------------#
class MAP:
    size = [[-300.0, -300.0], [300.0, 300.0]] # x, y最小值; x, y最大值
    start_pos = [-250, -250]                   # 起点坐标（左下角）
    end_pos = [250, 250]                    # 终点坐标（右上角）
    start_yaw = math.pi/4                 # 起点航向（朝向东北方向）
    obstacles = [                         # 障碍物, 要求为 geo.Polygon 或 带buffer的 geo.Point/geo.LineString
        geo.Point(0, 0).buffer(30),  # 减小障碍物半径
    ]

    @classmethod
    def show(cls):
        plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        plt.close('all')
        fig, ax = plt.subplots()
        ax.clear()
        cls.plot(ax)
        ax.scatter(cls.start_pos[0], cls.start_pos[1], s=30, c='k', marker='x', label='起点')
        ax.scatter(cls.end_pos[0], cls.end_pos[1], s=30, c='k', marker='o', label='终点')
        ax.legend(loc='best').set_draggable(True)
        plt.show(block=True)

    @classmethod
    def plot(cls, ax, title='Map'):
        ax.clear()
        ax.set_aspect('equal')
        ax.set_title(title)
        ax.set_xlabel("x")
        ax.set_ylabel("z")
        ax.grid(alpha=0.3, ls=':')
        ax.set_xlim(cls.size[0][0], cls.size[1][0])
        ax.set_ylim(cls.size[0][1], cls.size[1][1])
        # ax.invert_yaxis()
        for o in cls.obstacles:
            plot_polygon(o, ax=ax, facecolor='w', edgecolor='k', add_points=False)


# ----------------------------- ↓↓↓↓↓ 动力学避障环境 ↓↓↓↓↓ ------------------------------#
class Logger:
    pass

# 运动速度设置
V_LOW = 8  # 降低最小速度
V_HIGH = 15  # 最大速度
V_CRUISE = 12  # 巡航速度
# 质心动力学状态设置
STATE_LOW = [MAP.size[0][0], MAP.size[0][1], V_LOW, -math.pi]  # x, z, V, ψ
STATE_HIGH = [MAP.size[1][0], MAP.size[1][1], V_HIGH, math.pi]  # x, z, V, ψ

# 控制设置
CTRL_LOW = [-0.3, -0.2]  # 切向过载 + 偏航角速率(单位rad/s)
CTRL_HIGH = [0.3, 0.2]   # 切向过载 + 偏航角速率(单位rad/s)

# 距离设置
D_SAFE = 15  # 增加碰撞半径
D_BUFF = 40  # 增加缓冲距离
D_GOAL = 30  # 增加目标区域范围

# 序列观测长度
TIME_STEP = 4


class DynamicPathPlanning(gym.Env):
    """从力学与控制的角度进行多机协同规划 (东北天坐标系)
    >>> dx/dt = V * cos(ψ)
    >>> dz/dt = V * sin(ψ)
    >>> dV/dt = a
    >>> dψ/dt = psi_dot
    >>> u = [a, psi_dot]
    """

    def __init__(self, num_agents=2, max_episode_steps=500, dt=0.1, normalize_observation=True, old_gym_style=True):
        """
        Args:
            num_agents (int): 智能体数量. 默认2.
            max_episode_steps (int): 最大仿真步数. 默认500.
            dt (float): 决策周期. 默认0.1.
            normalize_observation (bool): 是否输出归一化的观测. 默认True.
            old_gym_style (bool): 是否采用老版gym接口. 默认True.
        """
        # 仿真
        self.dt = dt
        self.max_episode_steps = max_episode_steps
        self.log = Logger()
        # 障碍 + 雷达
        self.obstacles = MAP.obstacles
        # 状态空间 + 控制空间
        self.state_space = spaces.Box(np.array(STATE_LOW), np.array(STATE_HIGH))
        self.control_space = spaces.Box(np.array(CTRL_LOW), np.array(CTRL_HIGH))
        # 观测空间 + 动作空间
        # points_space = spaces.Box(-1, SCAN_RANGE, (TIME_STEP, SCAN_NUM, )) # seq_len, dim
        # vector_space = spaces.Box(np.array([OBS_STATE_LOW]*TIME_STEP), np.array([OBS_STATE_HIGH]*TIME_STEP)) # seq_len, dim
        # self.observation_space = spaces.Dict({'seq_points': points_space, 'seq_vector': vector_space})
        self.observation_space = spaces.Box(
            low=np.array([-1, -1, -1, -1, -1, -1, -1]),  # 7维观测空间
            high=np.array([1, 1, 1, 1, 1, 1, 1]),
            dtype=np.float32
        )
        self.action_space = spaces.Box(-1, 1, (len(CTRL_LOW), ))
        # 序列观测
        self.deque_points = deque(maxlen=TIME_STEP)
        self.deque_vector = deque(maxlen=TIME_STEP)
        # 环境控制
        self.__render_not_called = True
        self.__need_reset = True
        self.__norm_observation = normalize_observation
        self.__old_gym = old_gym_style
        # plt设置
        plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
        plt.rcParams['axes.unicode_minus'] = False
        plt.close("all")

    def reset(self, mode=0):
        """重置环境
           mode=0, 随机初始化起点终点, 速度、方向随机
           mode=1, 初始化起点终点到地图设置, 速度、方向随机
        """
        self.__need_reset = False
        self.time_step = 0
        # 初始化航程/状态/控制
        while 1:
            self.state = self.state_space.sample() # 随机初始化状态
            if mode == 0:
                self.center=np.array([random.randint(-200,-100),random.randint(-200,-100)], dtype=np.float32)
                self.theta=np.random.uniform(0, math.pi/2)
                self.state[:2] = self.center
                self.state[3] = self.theta
                self.start_pos = deepcopy(self.state[:2]) # 分不清引用传递, 暴力deepcopy就完事了
                goal_center = random.randint(100,200),random.randint(100,200)
                self.end_pos = np.array(goal_center, dtype=np.float32)
            else:
                self.start_pos = np.array(MAP.start_pos, dtype=np.float32)
                self.end_pos = np.array(MAP.end_pos, dtype=np.float32)
                self.state = np.array([*self.start_pos[:2], *self.state[2:]], dtype=np.float32)
            for o in self.obstacles:
                if o.contains(geo.Point(*self.start_pos)) \
                or o.contains(geo.Point(*self.end_pos)):
                    if mode != 0: raise ValueError("地图的初始/目标位置不能设置在障碍里面!!!")
                    break
            else:
                break
        self.L = 0.0                                                         # 航程
        self.ctrl = np.zeros(self.action_space.shape, dtype=np.float32) # 初始控制量
        # 初始化观测
        obs = self._get_obs(self.state)
        # 初始化记忆
        self.D_init = np.linalg.norm(self.end_pos - self.start_pos) # 初始到目标距离
        self.D_last = self.D_init # 上一时刻到目标距离
        # 重置log
        self.log.start_pos = self.start_pos       # 起点
        self.log.end_pos = self.end_pos           # 目标
        self.log.path = [self.start_pos]          # 路径
        self.log.ctrl = [self.ctrl]               # 控制
        self.log.speed = [self.state[2]]          # 速度
        self.log.yaw = [self.state[3]]            # 偏角
        self.log.length = [[self.L, self.D_last]] # 航程+距离
        # 输出
        if self.__old_gym:
            return obs
        return obs, {}

    def _get_ctrl(self, act, tau=0.9):
        """获取控制"""
        lb = self.control_space.low
        ub = self.control_space.high
        u = lb + (act + 1.0) * 0.5 * (ub - lb) # [-1,1] -> [lb,ub]
        u = np.clip(u, lb, ub) # NOTE 浮点数误差有时会出现类似 1.0000001 情况
        # smooth control signal
        if tau is not None:
            return (1.0 - tau) * self.ctrl + tau * u
        return u

    def _get_obs(self, state):
        """获取归一化的观测
        返回7维观测：
        - 自身位置(x, y)
        - 速度(V)
        - 偏航角(ψ)
        - 目标位置(end_x, end_y)
        - 障碍物标志(是否在障碍物附近)
        """
        x, y, V, ψ = state
        
        # 计算与障碍物的距离
        if self.obstacles is not None:
            dis_obs = self.obstacles[0].distance(geo.Point(x, y))
            obstacle_flag = 1 if dis_obs < D_SAFE else 0
        else:
            obstacle_flag = 0
            
        # 归一化观测
        obs = np.array([
            x / MAP.size[1][0],  # 归一化x坐标
            y / MAP.size[1][1],  # 归一化y坐标
            (V - V_LOW) / (V_HIGH - V_LOW),  # 归一化速度
            ψ / math.pi,  # 归一化航向角
            self.end_pos[0] / MAP.size[1][0],  # 归一化目标x坐标
            self.end_pos[1] / MAP.size[1][1],  # 归一化目标y坐标
            obstacle_flag  # 障碍物标志
        ], dtype=np.float32)
        
        return obs

    def _get_rew(self):
        """计算奖励
        奖励组成：
        1. 存活奖励：鼓励智能体持续探索
        2. 目标奖励：鼓励接近并到达目标
        3. 障碍物奖励：避免碰撞
        4. 航向奖励：保持朝向目标
        """
        x, y, V, ψ = self.state
        
        # 计算与目标的距离
        dis_goal = np.linalg.norm(self.end_pos - self.state[:2])
        distance_improvement = self.D_last - dis_goal
        self.D_last = dis_goal

        # 1. 存活奖励 - 给予小的正向激励
        survive_r = 0.1
        
        # 2. 目标奖励
        if dis_goal < D_GOAL:
            # 到达目标
            goal_r = 100.0
            done = True
            info = {'state': 'success'}
        else:
            # 未到达目标
            # 使用归一化的距离作为奖励基础
            normalized_dist = dis_goal / self.D_init
            goal_r = 2.0 * (1.0 - normalized_dist)  # 线性奖励
            done = False
            info = {'state': 'none'}

        # 3.障碍物奖励
        obstacle_r = 0
        if self.obstacles is not None:
            dis_obs = self.obstacles[0].distance(geo.Point(x, y))
            if dis_obs < D_SAFE:
                # 碰撞惩罚
                obstacle_r = -50.0
                done = True
                info = {'state': 'fail'}
            elif dis_obs < D_BUFF:
                # 安全距离惩罚，使用线性插值
                obstacle_r = -2.0 * (D_BUFF - dis_obs) / (D_BUFF - D_SAFE)

        # 4.航向奖励 - 鼓励朝向目标
        target_angle = math.atan2(self.end_pos[1] - y, self.end_pos[0] - x)
        angle_diff = abs(self._limit_angle(target_angle - ψ))
        # 使用余弦函数，并根据距离调整权重
        heading_r = 1.0 * np.cos(angle_diff)

        # 5.距离改善奖励 - 修改为非线性奖励
        if distance_improvement > 0:
            progress_r = 2.0 * distance_improvement / self.D_init
        else:
            progress_r = 4.0 * distance_improvement / self.D_init

        # 总奖励 - 使用较小的尺度
        reward = survive_r + goal_r + obstacle_r + heading_r + progress_r

        # 奖励缩放和裁剪
        reward = np.clip(reward, -50.0, 100.0)

        return reward, done, 0, obstacle_r, goal_r, info

    def step(self, act: np.ndarray, tau: float = None):
        """状态转移
        Args:
            act (np.ndarray): 动作a(取值-1~1).
            tau (float): 控制量u(取值u_min~u_max)的平滑系数: u = tau*u + (1-tau)*u_last. 默认None不平滑.
        """
        assert not self.__need_reset, "调用step前必须先reset"
        # 数值鸡分
        self.time_step += 1
        u = self._get_ctrl(act, tau)
        new_state = self._ode45(self.state, u, self.dt)
        truncated = False
        if self.time_step >= self.max_episode_steps:
            truncated = True
        elif new_state[0] > self.state_space.high[0] \
            or new_state[1] > self.state_space.high[1] \
            or new_state[0] < self.state_space.low[0] \
            or new_state[1] < self.state_space.low[1]:
            truncated = True
        # 更新航程/状态/唱跳rap篮球
        self.L += np.linalg.norm(new_state[:2] - self.state[:2])
        self.state = deepcopy(new_state)
        self.ctrl = deepcopy(u)
        # 获取转移元组
        obs = self._get_obs(new_state)
        rew, done, edge_r, obstacle_r, goal_r, info = self._get_rew()
        info["done"] = done
        info["truncated"] = truncated
        if truncated or done:
            info["terminal"] = True
            self.__need_reset = True
        else:
            info["terminal"] = False
        info["reward"] = rew
        info["edge_r"] = edge_r
        info["obstacle_r"] = obstacle_r
        info["goal_r"] = goal_r
        info["time_step"] = self.time_step
        info["voyage"] = self.L
        info["distance"] = self.D_last
        # 记录
        self.log.path.append(self.state[:2])
        self.log.ctrl.append(u)
        self.log.speed.append(self.state[2])
        self.log.yaw.append(self.state[3])
        self.log.length.append([self.L, self.D_last])
        # 输出
        if self.__old_gym:
            return obs, rew, done, edge_r, obstacle_r, goal_r, info
        return obs, rew, done, edge_r, obstacle_r, goal_r, truncated, info

    def _norm_obs(self, obs):
        """归一化观测"""
        if not self.__norm_observation:
            return obs
        obs['seq_vector'] = self._linear_mapping(
            obs['seq_vector'], 
            self.observation_space['seq_vector'].low, 
            self.observation_space['seq_vector'].high
        )
        obs['seq_points'] = self._normalize_points(obs['seq_points'])
        return obs

    def render(self, mode="human", figsize=[8,8]):
        """测试时可视化环境, 和step交替调用 (不要和plot一起调用, 容易卡)
        Args:
            mode (str): 显示模式. 默认"human".
            figsize (list): 图形大小. 默认[8,8].
        """
        assert not self.__need_reset, "调用render前必须先reset"
        
        # 创建绘图窗口
        if self.__render_not_called:
            self.__render_not_called = False
            plt.ion()  # 打开交互模式
            self.__fig = plt.figure("render", figsize=figsize)
            self.__ax = self.__fig.add_axes([0.1, 0.1, 0.8, 0.8])
            MAP.plot(self.__ax, "Path Plan Environment")
            
            # 创建轨迹线
            self.__plt_car_path, = self.__ax.plot([], [], 'b-', alpha=0.5, linewidth=1.5, label='Path')
            
            # 创建智能体当前位置点
            self.__plt_car_point = self.__ax.scatter([], [], s=50, c='b', marker='o', 
                                            label='Agent', edgecolor='k', linewidth=1)
            
            # 创建目标范围圆
            self.__plt_targ_range, = self.__ax.plot([], [], 'g:', linewidth=1.0, label='Target Range')
            
            # 创建目标点
            self.__plt_targ_point = self.__ax.scatter([], [], s=50, c='g', marker='*',
                                             label='Target', edgecolor='k', linewidth=1)
            
            # 添加图例
            self.__ax.legend(loc='best', bbox_to_anchor=(1.15, 1)).set_draggable(True)
            
            # 添加网格
            self.__ax.grid(True, linestyle='--', alpha=0.3)
            
            # 设置坐标轴标签
            self.__ax.set_xlabel('X Position')
            self.__ax.set_ylabel('Y Position')
            
            # 调整布局
            self.__fig.subplots_adjust(right=0.85)  # 为图例留出空间
            
            # 显示图形
            plt.show(block=False)
            plt.pause(0.1)  # 给窗口一些时间来显示
            
        
        # 更新绘图
        # 更新轨迹
        path_data = np.array(self.log.path)
        self.__plt_car_path.set_data(path_data[:, 0], path_data[:, 1])
        
        # 更新智能体位置
        current_pos = self.log.path[-1]
        self.__plt_car_point.set_offsets([current_pos])
        
        # 更新目标范围
        θ = np.linspace(0, 2*np.pi, 36)  # 增加点数使圆更平滑
        target_x = self.log.end_pos[0] + D_GOAL * np.cos(θ)
        target_y = self.log.end_pos[1] + D_GOAL * np.sin(θ)
        self.__plt_targ_range.set_data(target_x, target_y)
        
        # 更新目标点
        self.__plt_targ_point.set_offsets([self.log.end_pos])
        
        # 自动调整坐标轴范围
        self.__ax.relim()
        self.__ax.autoscale_view()
        
        # 刷新图形
        self.__fig.canvas.draw()
        self.__fig.canvas.flush_events()
        
        # 短暂暂停以允许图形更新
        plt.pause(0.001)

    def close(self): 
        """关闭环境"""
        self.__render_not_called = True
        self.__need_reset = True
        plt.close("render")

    def plot(self, file, figsize=[10,10], dpi=100):
        """训练时观察输出状态 (不要和render一起调用, 容易卡)"""
        file = Path(file).with_suffix(".png")
        file.parents[0].mkdir(parents=True, exist_ok=True)
        fig = plt.figure("Output", figsize=figsize)
        gs = fig.add_gridspec(2, 2) 
        ax1 = fig.add_subplot(gs[0, 0]) 
        ax2 = fig.add_subplot(gs[0, 1])
        ax3 = fig.add_subplot(gs[1, 0])
        ax4 = fig.add_subplot(gs[1, 1])
        MAP.plot(ax1, "Trajectory")
        ax1.scatter(*self.log.path[0], s=30, c='k', marker='x', label='start')
        ax1.scatter(*self.log.end_pos, s=30, c='k', marker='o', label='target')
        ax1.plot(*np.array(self.log.path).T, color='b', label='path') # [xxxxyyyy]
        ax1.legend(loc="best").set_draggable(True)
        ax2.set_title("Control Signal")
        ax2.set_xlabel("time step")
        ax2.set_ylabel("control")
        ctrl = np.array(self.log.ctrl).T
        for i, u in enumerate(ctrl):
            ax2.plot(u, label=f'u{i}')
        ax2.legend(loc="best").set_draggable(True)
        ax3.set_title("Speed Signal")
        ax3.set_xlabel("time step")
        ax3.set_ylabel("speed")
        ax3.plot(self.log.speed, label='V')
        ax3.legend(loc="best").set_draggable(True)
        ax4.set_title("Length Signal")
        ax4.set_xlabel("time step")
        ax4.set_ylabel("length")
        length = np.array(self.log.length).T
        ax4.plot(length[0], label='voyage')
        ax4.plot(length[1], label='distance')
        ax4.legend(loc="best").set_draggable(True)
        plt.tight_layout()
        fig.savefig(fname=file, dpi=dpi)
        plt.close("Output")

    @staticmethod
    def _limit_angle(x, domain=1):
        """限制角度 x 的区间: 1限制在(-π, π], 2限制在[0, 2π)"""
        x = x - x//(2*math.pi) * 2*math.pi # any -> [0, 2π)
        if domain == 1 and x > math.pi:
            return x - 2*math.pi           # [0, 2π) -> (-π, π]
        return x

    @staticmethod
    def _linear_mapping(x, x_min, x_max, left=0.0, right=1.0):  
        """x 线性变换: [x_min, x_max] -> [left, right]"""
        y = left + (right - left) / (x_max - x_min) * (x - x_min)
        return y

    @staticmethod
    def _vector_angle(x_vec, y_vec, EPS=1e-8):
        """计算向量 x_vec 与 y_vec 之间的夹角 [0, π]"""
        x = np.linalg.norm(x_vec) * np.linalg.norm(y_vec)
        y = np.dot(x_vec, y_vec)
        if x < EPS: # 0向量情况
            return 0.0
        if y < EPS: # 90°情况
            return math.pi/2
        return math.acos(np.clip(y/x, -1, 1)) # note: x很小的时候, 可能会超过+-1

    @staticmethod
    def _compute_azimuth(pos1, pos2, use_3d_pos=False):
        """东北天坐标系计算pos2相对pos1的方位角 [-π, π] 和高度角(3D情况) [-π/2, π/2] """
        if use_3d_pos:
            x, y, z = np.array(pos2) - pos1
            q = math.atan(z / (math.sqrt(x**2 + y**2) + 1e-8))  # 高度角 [-π/2, π/2]
            ε = math.atan2(y, x)  # 方位角 [-π, π]
            return ε, q
        else:
            x, y = np.array(pos2) - pos1
            return math.atan2(y, x)  # 方位角 [-π, π]

    @staticmethod
    def _fixed_wing_2d(s, t, u):
        """东北天坐标系平面运动ode模型
        s = [x, y, V, ψ]
        u = [a, psi_dot]
        """
        _, _, V, ψ = s
        a, psi_dot = u
        dsdt = [
            V * math.cos(ψ),  # x方向速度
            V * math.sin(ψ),  # y方向速度
            a,               # 加速度
            psi_dot         # 偏航角速率
        ]
        return dsdt

    @classmethod
    def _ode45(cls, s_old, u, dt):
        """微分方程积分"""
        s_new = odeint(cls._fixed_wing_2d, s_old, (0.0, dt), args=(u, )) # shape=(len(t), len(s))
        x, y, V, ψ = s_new[-1]
        V = np.clip(V, V_LOW, V_HIGH)
        ψ = cls._limit_angle(ψ)
        return np.array([x, y, V, ψ], dtype=np.float32) # deepcopy
